# build_tryagain.py
# usage: results.jsonl 没有生成问题清单，重新写回 questions 文件
# python ./src/tools/build_tryagain.py --test "./data/test.jsonl" --results "./result/results.jsonl" --questions "./data/questions.jsonl"

import argparse
import json
import sys

def load_jsonl_file(filepath):
    """
    从指定文件路径加载 JSONL（每行一个 JSON 对象）文件。
    返回一个以 'id' 为键的字典。
    """
    data = {}
    try:
        with open(filepath, 'r', encoding='utf-8') as file:
            for line_num, line in enumerate(file, start=1):
                line = line.strip()
                if not line:  # 跳过空行
                    continue
                try:
                    obj = json.loads(line)
                    if 'id' not in obj:
                        print(f"警告: 第 {line_num} 行缺少 'id' 字段，已跳过。")
                        continue
                    data[obj['id']] = obj
                except json.JSONDecodeError as e:
                    print(f"错误: 第 {line_num} 行 JSON 解析失败: {e}")
                    sys.exit(1)
    except FileNotFoundError:
        print(f"错误: 文件未找到: {filepath}")
        sys.exit(1)
    except Exception as e:
        print(f"错误: 读取文件时发生错误: {e}")
        sys.exit(1)
    return data

def main():
    parser = argparse.ArgumentParser(description="从 --test 文件中提取在 --results 文件中不存在的记录，并写入 --questions 文件。")
    parser.add_argument('--test', required=True, help='包含所有测试问题的 JSONL 文件路径')
    parser.add_argument('--results', required=True, help='包含已处理结果的 JSONL 文件路径')
    parser.add_argument('--questions', required=True, help='输出缺失问题的 JSONL 文件路径')

    args = parser.parse_args()

    # 加载文件数据
    print(f"正在加载 --test 文件: {args.test}")
    test_data = load_jsonl_file(args.test)

    print(f"正在加载 --results 文件: {args.results}")
    results_data = load_jsonl_file(args.results)

    # 找出在 test 中但不在 results 中的记录
    missing_ids = set(test_data.keys()) - set(results_data.keys())
    missing_records = [test_data[id] for id in sorted(missing_ids)]  # 按 ID 排序输出

    # 写入输出文件
    try:
        with open(args.questions, 'w', encoding='utf-8') as outfile:
            for record in missing_records:
                outfile.write(json.dumps(record, ensure_ascii=False) + '\n')
        print(f"成功: 找到 {len(missing_records)} 条缺失记录，已写入 '{args.questions}'。")
    except Exception as e:
        print(f"错误: 写入文件 '{args.questions}' 时发生错误: {e}")
        sys.exit(1)

if __name__ == '__main__':
    main()